package com.ruoyi.znjc.ai.service;

import com.ruoyi.znjc.domain.ZnjcMessages;
import com.ruoyi.znjc.service.IZnjcMessagesService;
import okhttp3.*;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.List;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

/**
 * ChatSseService
 * 通过百度智能云千帆 API 实现基于 SSE 的流式聊天服务（支持知识库优先回答）
 */
@Service
public class ChatSseService {

    @Autowired
    private IZnjcMessagesService znjcMessagesService;// 查询历史消息

    private final ExecutorService executor = Executors.newCachedThreadPool();
    // 替换为你自己的百度智能云千帆API Key
    private static final String API_KEY = "Bearer bce-v3/ALTAK-hhFkTe3pihmjnPxLjMlcW/4d8405b990fbfd6e81eddbc401681b14a2133e32";
    // 千帆模型接口 URL
    private static final String API_URL = "https://qianfan.baidubce.com/v2/chat/completions";
    // 知识库接口URL
    private static final String KB_SEARCH_URL = "https://qianfan.baidubce.com/v2/knowledgebases/query";
    // 你的知识库ID，必填
    private static final String KNOWLEDGE_BASE_ID = "f025c9c9-1ffa-4c60-850c-060d63dd719d";

    public SseEmitter chat(String question, Integer conversationId) {
        SseEmitter emitter = new SseEmitter(0L);

        executor.submit(() -> {
            OkHttpClient client = new OkHttpClient.Builder()
                    .readTimeout(0, java.util.concurrent.TimeUnit.SECONDS)
                    .build();

            MediaType mediaType = MediaType.parse("application/json");

            // 1. 调用知识库检索
            String kbContent = queryKnowledgeBase(question);

            // 2. 查询历史消息
            List<ZnjcMessages> historyMessages = znjcMessagesService.selectMessagesByConversationId(Long.valueOf(conversationId));

            // 3. 构造 messages JSON 数组
            StringBuilder messagesBuilder = new StringBuilder();
            messagesBuilder.append("[\n");

            // 如果知识库有内容，放入system消息，提示严格输出知识库内容
            if (kbContent != null && !kbContent.trim().isEmpty()) {
                String systemContent = "你是一个知识库回答生成器。\n" +
                        "以下是知识库内容，请严格逐字逐句输出，不得省略或改写：\n" + kbContent;
                messagesBuilder.append("  {\"role\": \"system\", \"content\": \"")
                        .append(escapeJson(systemContent))
                        .append("\"},\n");
            }

            // 历史消息
            for (int i = 0; i < historyMessages.size(); i++) {
                ZnjcMessages msg = historyMessages.get(i);
                String roleStr = (msg.getRole() == 0) ? "assistant" : "user";
                messagesBuilder.append("  {\"role\": \"")
                        .append(roleStr)
                        .append("\", \"content\": \"")
                        .append(escapeJson(msg.getContent()))
                        .append("\"},\n");
            }

            // 当前用户问题
            messagesBuilder.append("  {\"role\": \"user\", \"content\": \"")
                    .append(escapeJson(question))
                    .append("\"\n}");

            messagesBuilder.append("\n]");

            // 4. 构造请求体JSON
            String bodyJson = "{\n" +
                    "  \"model\": \"ernie-3.5-8k\",\n" +
                    "  \"stream\": true,\n" +
                    "  \"messages\": " + messagesBuilder.toString() + "\n" +
                    "}";

            RequestBody body = RequestBody.create(mediaType, bodyJson);

            Request request = new Request.Builder()
                    .url(API_URL)
                    .post(body)
                    .addHeader("Authorization", API_KEY)
                    .addHeader("Content-Type", "application/json")
                    .build();

            try (Response response = client.newCall(request).execute()) {
                if (!response.isSuccessful()) {
                    String errMsg = "请求失败: " + response.code() + " - " + response.message();
                    emitter.send(SseEmitter.event().name("error").data(errMsg));
                    emitter.complete();
                    return;
                }

                BufferedReader reader = new BufferedReader(new InputStreamReader(response.body().byteStream()));
                String line;
                while ((line = reader.readLine()) != null) {
                    if (line.startsWith("data: ")) {
                        String data = line.substring("data: ".length());
                        if ("[DONE]".equals(data)) {
                            emitter.send(SseEmitter.event().name("done").data("[DONE]"));
                            break;
                        }
                        emitter.send(SseEmitter.event().data(data));
                    }
                }
                emitter.complete();
            } catch (IOException e) {
                try {
                    emitter.send(SseEmitter.event().name("error").data("后端异常：" + e.getMessage()));
                } catch (IOException ex) {
                    ex.printStackTrace();
                }
                emitter.completeWithError(e);
            }
        });

        return emitter;
    }

    /**
     * 调用知识库检索接口，返回命中的内容字符串，没有命中返回null
     */
    private String queryKnowledgeBase(String question) {
        OkHttpClient client = new OkHttpClient();
        MediaType mediaType = MediaType.parse("application/json");
        String bodyJson = "{"
                + "\"knowledge_base_id\": \"" + KNOWLEDGE_BASE_ID + "\","
                + "\"query\": \"" + escapeJson(question) + "\","
                + "\"top_k\": 1"
                + "}";
        RequestBody body = RequestBody.create(mediaType, bodyJson);
        Request request = new Request.Builder()
                .url(KB_SEARCH_URL)
                .post(body)
                .addHeader("Authorization", API_KEY)
                .addHeader("Content-Type", "application/json")
                .build();
        try (Response response = client.newCall(request).execute()) {
            if (response.isSuccessful() && response.body() != null) {
                String json = response.body().string();
                return extractContent(json);
            }
        } catch (IOException e) {
            e.printStackTrace();
        }
        return null;
    }

    /**
     * 简单示例，解析返回JSON中第一个命中文本的content字段，建议用JSON库替换
     */
    private String extractContent(String json) {
        // TODO: 根据实际返回的JSON结构修改解析逻辑
        int idx = json.indexOf("\"content\":\"");
        if (idx != -1) {
            idx += 10;
            int end = json.indexOf("\"", idx);
            if (end > idx) {
                String content = json.substring(idx, end);
                // JSON中转义的换行符等转换回正常字符
                content = content.replace("\\n", "\n").replace("\\\"", "\"").replace("\\\\", "\\");
                return content;
            }
        }
        return null;
    }

    /**
     * 转义字符串，防止JSON格式错误
     */
    private String escapeJson(String text) {
        if (text == null) return "";
        return text.replace("\\", "\\\\")
                .replace("\"", "\\\"")
                .replace("\n", "\\n")
                .replace("\r", "\\r");
    }
}